/* 
* Copyright (C) 2010, John Asmuth

 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
 You may obtain a copy of the License at

     http://www.apache.org/licenses/LICENSE-2.0

 Unless required by applicable law or agreed to in writing, software
 distributed under the License is distributed on an "AS IS" BASIS,
 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
* 
*  $Revision$
*  $Date$
*  $Author$
*  $HeadURL$
* 
*/

package rlglue

import (
	"fmt"
	"math/rand"
	"strconv"
	"strings"
)

type Environment interface {
	EnvInit() (taskString string)
	EnvStart() (obs Observation)
	EnvStep(action Action) (obs Observation, r float64, t bool)
	EnvCleanup()
	EnvMessage(message string) (reply string)
}

type DefaultEnvironment struct{}

func (*DefaultEnvironment) EnvInit() (taskString string)                        { return }
func (*DefaultEnvironment) EnvStart() (obs Observation)                         { return }
func (*DefaultEnvironment) EnvStep(Action) (obs Observation, r float64, t bool) { return }
func (*DefaultEnvironment) EnvCleanup()                                         {}
func (*DefaultEnvironment) EnvMessage(message string) (reply string) {
	if strings.HasPrefix(message, "seed ") {
		seedStr := message[len("seed "):]
		if seed, err := strconv.ParseInt(seedStr, 10, 64); err == nil {
			rand.Seed(seed)
		}
	}
	return
}

func LoadEnvironment(env Environment) (err error) {
	ce := NewClientEnvironment(env)
	if err = ce.AutoConnect(); err != nil {
		return
	}
	if err = ce.RunEnvironmentEventLoop(); err != nil {
		return
	}
	err = ce.Close()
	return
}

type ClientEnvironment struct {
	env Environment
	nw  *network
}

func NewClientEnvironment(env Environment) *ClientEnvironment {
	return &ClientEnvironment{env, nil}
}

func (cl *ClientEnvironment) onEnvInit() (err error) {
	taskSpec := cl.env.EnvInit()
	cl.nw.ClearSendBuffer()
	cl.nw.PutInt(kEnvInit)
	cl.nw.PutInt(int32(len(taskSpec)) + kIntSize)
	cl.nw.PutString(taskSpec)
	return
}

func (cl *ClientEnvironment) onEnvStart() (err error) {
	obs := cl.env.EnvStart()
	size := cl.nw.SizeOfObservation(obs)
	cl.nw.ClearSendBuffer()
	cl.nw.PutInt(kEnvStart)
	cl.nw.PutInt(size)
	cl.nw.PutObservation(obs)
	return
}

func (cl *ClientEnvironment) onEnvStep() (err error) {
	action, err := cl.nw.GetAction()
	if err != nil {
		return
	}
	obs, r, t := cl.env.EnvStep(action)
	size := cl.nw.SizeOfRewardObservation(obs)
	cl.nw.ClearSendBuffer()
	cl.nw.PutInt(kEnvStep)
	cl.nw.PutInt(size)
	cl.nw.PutRewardObservation(obs, r, t)
	return
}

func (cl *ClientEnvironment) onEnvCleanup() (err error) {
	cl.env.EnvCleanup()
	cl.nw.ClearSendBuffer()
	cl.nw.PutInt(kEnvCleanup)
	cl.nw.PutInt(0)
	return
}

func (cl *ClientEnvironment) onEnvMessage() (err error) {
	message, err := cl.nw.GetString()
	if err != nil {
		return
	}
	reply := cl.env.EnvMessage(message)
	cl.nw.ClearSendBuffer()
	cl.nw.PutInt(kEnvMessage)
	cl.nw.PutInt(int32(len(reply)) + kIntSize)
	cl.nw.PutString(reply)
	return
}

func (cl *ClientEnvironment) Connect(host string, port int) (err error) {
	cl.nw = new(network)
	err = cl.nw.Connect(host, port)
	if err != nil {
		return
	}
	cl.nw.ClearSendBuffer()
	cl.nw.PutInt(kEnvironmentConnection)
	cl.nw.PutInt(0)
	err = cl.nw.Send()
	return
}

func (ce *ClientEnvironment) AutoConnect() (err error) {
	host, port := GetEnvHostPort()
	err = ce.Connect(host, port)
	return
}

func (cl *ClientEnvironment) Close() error {
	return cl.nw.Close()
}

func (cl *ClientEnvironment) RunEnvironmentEventLoop() (err error) {
	var envState int32 = 0

	for envState != kRLTerm {
		cl.nw.ClearRecvBuffer()
		// We may have received the header and part of the payload
		// We need to keep track of how much of the payload was recv'd
		var recvSize int32
		recvSize, err = cl.nw.Recv(2 * kIntSize)
		if err != nil {
			return
		}
		recvSize -= 2 * kIntSize

		envState, err = cl.nw.GetInt()
		if err != nil {
			return
		}

		//fmt.Printf("env state: %d\n", envState)

		//print("env getting dataSize\n")
		var dataSize int32
		dataSize, err = cl.nw.GetInt()
		if err != nil {
			return
		}
		//print("env done\n")

		remaining := dataSize - recvSize
		if remaining < 0 {
			//print("Remaining was less than 0!\n")
			remaining = 0
		}

		//fmt.Printf("env reading remaining %d\n", remaining)
		_, err = cl.nw.Recv(remaining)
		if err != nil {
			return
		}
		//print("env done\n")

		// Already read the header, discard it
		//cl.nw.GetInt()
		//cl.nw.GetInt()

		//print("env calling method\n")

		switch envState {
		case kEnvInit:
			cl.onEnvInit()
		case kEnvStart:
			cl.onEnvStart()
		case kEnvStep:
			cl.onEnvStep()
		case kEnvCleanup:
			cl.onEnvCleanup()
		case kEnvMessage:
			cl.onEnvMessage()
		case kRLTerm:
			break
		default:
			panic(fmt.Sprintf("Unknown message type: %d", envState))
		}

		err = cl.nw.Send()
		if err != nil {
			return
		}
	}
	return
}
